package com.hucsmn.xcopy;

import java.io.*;
import java.nio.file.AccessDeniedException;
import java.nio.file.FileAlreadyExistsException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;

/**
 * 文件复制线程池。
 */
public class MultiThreadCopy {
    /**
     * 每个文件的默认缓冲区大小。
     */
    public static final int BUFFER_SIZE = 8192;

    /**
     * 默认线程数。
     */
    public static final int THREADS = 2 * Runtime.getRuntime().availableProcessors();

    private File src;
    private File dst;
    private FileTree srcTree;

    private int bufferSize;
    private int threads;
    private ExecutorService pool = null;

    /**
     * 创建文件复制线程池。
     *
     * @param src 源文件或目录。
     * @param dst 目标文件或目录。
     * @throws IOException 源文件夹读取失败。
     */
    public MultiThreadCopy(File src, File dst) throws IOException {
        this(src, dst, THREADS);
    }

    /**
     * 创建文件复制线程池。
     *
     * @param src 源文件或目录。
     * @param dst 目标文件或目录。
     * @param threads 指定线程数量。
     * @throws IOException 源文件夹读取失败。
     */
    public MultiThreadCopy(File src, File dst, int threads) throws IOException {
        this(src, dst, threads, BUFFER_SIZE);
    }

    /**
     * 创建文件复制线程池。
     *
     * @param src 源文件或目录。
     * @param dst 目标文件或目录。
     * @param threads 指定线程数量。
     * @param bufferSize 每个复制任务使用的缓冲区大小。
     * @throws IOException 源文件夹读取失败。
     */
    public MultiThreadCopy(File src, File dst, int threads, int bufferSize) throws IOException {
        setSrc(src);
        setDst(dst);
        setThreads(threads);
        setBufferSize(bufferSize);
    }

    /**
     * 启动多线程复制。
     *
     * @throws IOException 任何一个文件夹无法创建的异常，或者所有文件复制任务的出错信息汇总。
     */
    public synchronized void start() throws IOException {
        // 避免相同位置得文件复制，避免创建文件时源文件被截断。
        if (src.equals(dst)) {
            return;
        }

        // 先保证文件夹已经创建，防止多线程复制时创建不了文件。
        for (File srcDir : srcTree.dirs()) {
            File dstDir = dstLocationOf(srcDir);
            if (!dstDir.exists()) {
                System.out.println("正在创建文件夹 " + dstDir.getPath());
                if (!dstDir.mkdir()) {
                    throw new AccessDeniedException(dstDir.getPath());
                }
            } else if (!dstDir.isDirectory()) {
                throw new FileAlreadyExistsException(dstDir.getPath());
            }
        }

        // 创建或复用线程池。
        if (pool == null) {
            pool = Executors.newFixedThreadPool(threads);
        }

        // 将全部的文件复制任务都添加进线程池内。
        List<File> files = srcTree.files();
        ArrayList<Future<Long>> tasks = new ArrayList<>(files.size());
        for (File srcFile : files) {
            File dstFile = dstLocationOf(srcFile);
            Future<Long> task = pool.submit(() -> copyFile(srcFile, dstFile, bufferSize));
            tasks.add(task);
        }

        // 等待全部文件复制完成，并收集任务的出错信息。
        List<Throwable> failures = new ArrayList<>();
        long total = 0;
        boolean done;
        do {
            done = true;
            for (int i = 0; i < tasks.size(); i++) {
                Future<Long> task = tasks.get(i);
                if (task == null) {
                    continue;
                }
                if (!task.isDone()) {
                    done = false;
                    continue;
                }

                try {
                    total += task.get();
                } catch (InterruptedException e) {
                    failures.add(e);
                } catch (ExecutionException e) {
                    failures.add(e.getCause());
                } finally {
                    tasks.set(i, null);
                }
            }
        } while (!done);

        if (failures.size() > 0) {
            throw new MultiThreadCopyException(failures.toString(), failures);
        } else {
            System.out.println("共复制了 " + total + " 字节的文件");
        }
    }

    /**
     * 使用缓冲复制文件。
     *
     * @param srcFile 源文件位置。
     * @param dstFile 目标文件位置。
     * @param bufferSize 缓冲区大小。
     * @throws IOException 源文件无法打开并读取、目标文件无法并创建写入。
     */
    private static long copyFile(File srcFile, File dstFile, int bufferSize) throws IOException {
        if (bufferSize <= 0) {
            throw new IllegalArgumentException();
        }

        System.out.println("正在将 " + srcFile.getPath() + " 复制到 " + dstFile.getPath() + " 位置...");
        long total = 0;
        try (InputStream in = new FileInputStream(srcFile);
             OutputStream out = new FileOutputStream(dstFile)) {
            byte[] buf = new byte[BUFFER_SIZE];
            int count = 0;
            while ((count = fillBuffer(in, buf)) >= 0) {
                out.write(buf, 0, count);
                total += count;
            }
        }
        return total;
    }

    /**
     * 尽可能地将更多的输入数据填充进缓冲区。
     *
     * @param in 输入流。
     * @param buf 缓冲区。
     * @return 缓冲区填入的字节数，遇到文件末尾并且没有读取到任何数据时，返回 -1。
     * @throws IOException 文件读取错误。
     */
    private static int fillBuffer(InputStream in, byte[] buf) throws IOException {
        int count = 0;
        while (count < buf.length) {
            int n = in.read(buf, count, buf.length - count);
            if (n < 0) {
                break;
            }
            count += n;
        }
        return count != 0 ? count : -1;
    }

    /**
     * 获取目标文件或文件夹位置。
     *
     * @param srcLocation 源文件或源文件夹位置。
     * @return 目标文件或文件夹位置。
     */
    private File dstLocationOf(File srcLocation) {
        String srcPrefix = src.getPath();
        String dstPrefix = dst.getPath();
        String srcPath = srcLocation.getPath();
        if (!srcPath.startsWith(srcPrefix)) {
            throw new IllegalArgumentException();
        }
        String dstPath = dstPrefix + srcPath.substring(srcPrefix.length());
        return new File(dstPath);
    }

    /**
     * 获取源文件或文件夹位置。
     *
     * @return 源文件或文件夹位置。
     */
    public File getSrc() {
        return src;
    }

    /**
     * 设置源文件或文件夹位置。
     *
     * @param src 源文件或文件夹位置。
     * @throws IOException 无法读取源文件文件夹。
     */
    public synchronized void setSrc(File src) throws IOException {
        if (src == null) {
            throw new NullPointerException();
        }
        this.src = src;
        this.srcTree = new FileTree(this.src);
    }

    /**
     * 获取目标文件或文件夹位置。
     *
     * @return 目标文件或文件夹位置。
     */
    public File getDst() {
        return dst;
    }

    /**
     * 设置目标文件或文件夹位置。
     *
     * @param dst 目标文件或文件夹位置。
     * @throws IOException 整理文件路径时出错。
     */
    public synchronized void setDst(File dst) throws IOException {
        if (dst == null) {
            throw new NullPointerException();
        }
        this.dst = dst.getCanonicalFile();
    }

    /**
     * 获取线程池大小。
     *
     * @return 线程池大小。
     */
    public int getThreads() {
        return threads;
    }

    /**
     * 设置线程池大小。
     *
     * @param threads 线程池大小。
     */
    public synchronized void setThreads(int threads) {
        if (threads <= 0) {
            throw new IllegalArgumentException();
        }
        this.threads = threads;
        resetThreadPool();
    }

    /**
     * 释放线程池。
     */
    public synchronized void resetThreadPool() {
        if (pool != null) {
            pool.shutdown();
            pool = null;
        }
    }

    /**
     * 获取每一个任务的缓冲区大小。
     *
     * @return 缓冲区大小。
     */
    public int getBufferSize() {
        return bufferSize;
    }

    /**
     * 设置每一个任务的缓冲区大小。
     *
     * @param bufferSize 缓冲区大小。
     */
    public synchronized void setBufferSize(int bufferSize) {
        if (bufferSize <= 0) {
            throw new IllegalArgumentException();
        }
        this.bufferSize = bufferSize;
    }
}
